import imgaug as ia
import imgaug.augmenters as iaa
from imgaug.augmentables.segmaps import SegmentationMapsOnImage
import imageio
import numpy as np
import os,glob
import natsort
import cv2
from utils import config
from PIL import Image

'random channel exchange'
def RandomBGRRGBInverse(image, label):
    if np.random.random() > .5:
      image = image[:, :, ::-1]
    return image, label
def withChannels(image, label):
  if np.random.random() >.5:
    aug = iaa.WithChannels()
    image, label = aug(image=image, segmentation_maps=label)
  return image, label


'随机遮挡，分割图无影响'
def croarsedropout_1(image, label):
  if np.random.random() >.7:
    aug = iaa.CoarseDropout(0.01, size_percent=(0.01,0.5))
    image, label = aug(image=image, segmentation_maps=label)

  return image, label


'随机遮挡+左右翻转，分割图也要翻转'
def croarsedropout_2(image, label):
  if np.random.random() > .8:
    aug = iaa.CoarseDropout(0.01, size_percent=0.3)
    image, label = aug(image=image, segmentation_maps=label)
    aug = iaa.Fliplr()
    image, label = aug(image=image, segmentation_maps=label)
  return image, label


'随机遮挡+上下翻转，分割图也要翻转'
def croarsedropout_3(image, label):
  if np.random.random() > .8:
    aug = iaa.CoarseDropout(0.01, size_percent=0.3)
    image, label = aug(image=image, segmentation_maps=label)
    aug = iaa.Flipud()
    image, label = aug(image=image, segmentation_maps=label)
  return image, label


'弹性形变，分割图也会同步形变'  # 尽量不用该增强方法，会产生黑边影响最终的效果
def elastic(image, label):
  if np.random.random() > .5:
    aug = iaa.ElasticTransformation(alpha=120, sigma=20)
    image, label = aug(image=image, segmentation_maps=label)
  return image, label

# def elastic1(image, label):
#     aug = iaa.ElasticTransformation(alpha=120, sigma=20)
#     image, label = aug(image=image, segmentation_maps=label)
#     aug = iaa.Fliplr()
#     img, mask = aug(image=image, segmentation_maps=label)
#     return img, mask


'水平翻转，分割图也会同步翻转'
def fliplevel(image, label):
  if np.random.random() > .7:
    aug = iaa.Fliplr()
    image, label = aug(image=image, segmentation_maps=label)
  return image, label


'上下翻转，分割图也会同步翻转'
def flipup(image, label):
  if np.random.random() > .7:
    aug = iaa.Flipud()
    image, label = aug(image=image, segmentation_maps=label)
  return image, label


'先水平翻转再上下翻转，分割图会同步翻转'
def flip(image, label):
  if np.random.random() > .5:
    auglr = iaa.Fliplr()
    img1, mask1 = auglr(image=image, segmentation_maps=label)
    augup = iaa.Flipud()
    image, label = augup(image=img1, segmentation_maps=mask1)
  return image, label


'只做明亮度增强'
def multiply_1(image, label):
  if np.random.random() > .7:
    aug = iaa.Multiply((0.5, 1.5),per_channel=True)
    image, label = aug(image=image, segmentation_maps=label)
  return image, label


'明亮度+随机遮挡，不影响分割图'
def multiply_2(image, label):
  if np.random.random() > .8:
    aug = iaa.Multiply((0.5, 1.5))
    image, label = aug(image=image, segmentation_maps=label)
    aug = iaa.CoarseDropout(0.01, size_percent=(0.01,0.04))
    image, label = aug(image=image, segmentation_maps=label)
  return image, label


'明暗度增强+上下翻转，分割图同步翻转'
def multiply_3(image, label):
  if np.random.random() > .5:
    aug = iaa.Multiply((0.1, 2.0))
    img1, mask1 = aug(image=image, segmentation_maps=label)
    augup = iaa.Flipud()
    image, label = augup(image=img1, segmentation_maps=mask1)
  return image, label


'只做高斯模糊'
def gaussianblur_1(image, label):
  if np.random.random() > .7:
    aug = iaa.GaussianBlur(sigma=(0, 3.0))
    image, label = aug(image=image, segmentation_maps=label)
  return image, label


'高斯模糊+随机遮挡'
def gaussianblur_2(image, label):
  aug = iaa.GaussianBlur(sigma=(0, 3.0))
  image, label = aug(image=image, segmentation_maps=label)
  aug = iaa.CoarseDropout(0.01, size_percent=0.3)
  img, mask = aug(image=image, segmentation_maps=label)
  return img, mask


'高斯模糊+水平翻转，分割图同步翻转'
def gaussianblur_3(image, label):
  aug = iaa.GaussianBlur(sigma=(0, 3.0))
  img1, mask1 = aug(image=image, segmentation_maps=label)
  auglr = iaa.Fliplr()
  img, mask = auglr(image=img1, segmentation_maps=mask1)
  return img, mask


def contrastNormalization(image,label):
  if np.random.random()>.7:
    aug=iaa.contrast.LinearContrast()
    image,label=aug(image=image,segmentation_maps=label)
  return image,label



def choice_(image_path, label_path, flag):
  print(image_path, label_path)
  image = imageio.imread(image_path)
  label = imageio.imread(label_path)
  #--------------------------------------------------------
  #               single image
  #-------------------------------------------------------
  # path1 = "/data2/enducation/datas/answer_card/unet_data/20211105/imgs/2021-10-16_02_07_44.jpg"
  # path2 = "/data2/enducation/datas/answer_card/unet_data/20211105/labels/2021-10-16_02_07_44.jpg"
  # image = Image.open(path1)
  # image = cv2.cvtColor(np.asanyarray(image), cv2.COLOR_RGB2BGR)
  # label = Image.open(path2)
  # label = cv2.cvtColor(np.asanyarray(label), cv2.COLOR_RGB2BGR)

  segmentation = SegmentationMapsOnImage(label, shape=image.shape)

  if flag == 0:
    return croarsedropout_1(image, segmentation)
  elif flag == 1:
    return croarsedropout_2(image, segmentation)
  elif flag == 2:
    return croarsedropout_3(image, segmentation)
  elif flag == 3:
    return elastic(image, segmentation)
  elif flag == 4:
    return fliplevel(image, segmentation)
  elif flag == 5:
    return flipup(image, segmentation)
  elif flag == 6:
    return multiply_1(image, segmentation)
  elif flag == 7:
    return multiply_2(image, segmentation)
  elif flag == 8:
    return multiply_3(image, segmentation)
  elif flag == 9:
    return gaussianblur_1(image, segmentation)
  elif flag == 10:
    return RandomBGRRGBInverse(image,segmentation)
  elif flag == 11:
    return contrastNormalization(image,segmentation)

def Random_imgaug(image, label, flag):
  segmentation = SegmentationMapsOnImage(label, shape=image.shape)
  if flag == 0:
    return croarsedropout_1(image, segmentation)
  elif flag == 1:
    return croarsedropout_2(image, segmentation)
  elif flag == 2:
    return croarsedropout_3(image, segmentation)
  elif flag == 3:
    return elastic(image, segmentation)
  elif flag == 4:
    return fliplevel(image, segmentation)
  elif flag == 5:
    return flipup(image, segmentation)
  elif flag == 6:
    return multiply_1(image, segmentation)
  elif flag == 7:
    return multiply_2(image, segmentation)
  elif flag == 8:
    return multiply_3(image, segmentation)
  elif flag == 9:
    return gaussianblur_1(image, segmentation)
  elif flag == 10:
    return RandomBGRRGBInverse(image,segmentation)
  elif flag == 11:
    return contrastNormalization(image,segmentation)


if __name__ == '__main__':
  # 该地址为更改名字后保存下来的图片地址（更改名字后的格式：x.jpg,x.png），在该文件中进行操作
  path = r'/data2/enducation/datas/answer_card/unet_data/20211109/'  #imgs  and labels
  # 该地址用来保存做了数据增强后的图片
  save_path = r'/data2/enducation/datas/answer_card/unet_data/20211109/'
  img_path_ = []
  mask_path_ = []
  data_count =0
  fan = []
  for file_name in os.listdir(path):
    if file_name =="imgs":
      for img_path in glob.glob(os.path.join(path,file_name,"*.png")):
        img_path_.append(img_path)
        data_count+=1
    else:
      for label_path in glob.glob(os.path.join(path,file_name,"*.png")):
        mask_path_.append(label_path)
  img_path = natsort.natsorted(img_path_)
  mask_path = natsort.natsorted(mask_path_)


  for i in range(data_count):
      image, mask = choice_(img_path[i],mask_path[i], np.random.randint(0, 1))
      # image = imageio.imread(img_path[i])
      # label = imageio.imread(mask_path[i])
      # segmentation = SegmentationMapsOnImage(label, shape=image.shape)
      # # aug = iaa.ElasticTransformation(alpha=120, sigma=20)
      # # image, label = aug(image=image, segmentation_maps=label)
      # seq=iaa.Sequential([iaa.ElasticTransformation(alpha=120, sigma=20)])
      # image, mask=seq(image=image,segmentation_maps=label)

      # ia.imshow(image)
      # ia.imshow(mask.get_arr())
      #-------------------------------
      cv2.imshow("image",image)
      # cv2.imshow("mask",mask.get_arr())
      cv2.waitKey(0)
      cv2.destroyAllWindows()

      # imageio.imwrite('epoch.jpg', image)
      # imageio.imwrite('epoch1.png', mask.get_arr())
  # exit()
